#!/usr/bin/env python
# -*- coding: utf8 -*-

# cache & lazy eval!
# DON'T cache things that may be changed by users except sysadmin

#import json
try:
    from django.utils import simplejson as json
except:
    import json
import urllib
import urllib2
import datetime
from cookielib import CookieJar
import time
import settings

def _formData(**data):
    for k in data:
        if isinstance(data[k], unicode):
            data[k] = data[k].encode('utf8')
    return urllib.urlencode(data)

class MediaWikiException(Exception):
    pass

class MediaWikiLocalException(MediaWikiException):
    pass

class MediaWikiApiError(MediaWikiException):
    def __init__(self, code=u'', info=u'', request={}):
        if isinstance(code, MediaWikiApiError):
            self.__init__(code.code, code.info, code.request)
        else:
            message = u'%s: %s; request: %s' % (code, info, request)
            super(MediaWikiApiError, self).__init__(message)
            self._code = code
            self._request = request
    code = property(lambda self: self._code)
    request = property(lambda self: self._request)

class PageNotExists(MediaWikiException):
    pass

class PageNotSaved(MediaWikiException):
    pass

class PageNotMoved(MediaWikiException):
    pass

class PageNotDeleted(MediaWikiException):
    pass

class SavedRevision(PageNotSaved):
    pass

class LoginFailed(MediaWikiException):
    pass

class MediaWikiCookieProcessor(urllib2.HTTPCookieProcessor):
    def http_response(self, request, response):
        cookies = response.info().getallmatchingheaders('Set-Cookie')
        if len(cookies) == 1:
            #GAE replace all '\r\n' from cookies, cause fail to login
            cookies = cookies[0].replace('HttpOnly', 'HttpOnly\r\n').replace('httponly\n', 'HttpOnly\r\n').replace('httponly', 'HttpOnly\r\n')
            cookies = ['Set-Cookie: %s\r\n' % ck.strip() for ck in cookies.split('\r\n,')]
            for h in response.info().headers:
                if h.strip().startswith('Set-Cookie'):
                    h = cookies[0]
            cookies.pop(0)
            response.info().headers.extend(cookies)
        self.cookiejar.extract_cookies(response, request)
        return response

def _fromTimestamp(ts):
    try:
        return datetime.datetime.strptime(ts, '%Y-%m-%dT%H:%M:%SZ')
    except Exception:
        return datetime.datetime.strptime(ts, '%Y%m%d%H%M%S')

def _namespace(namespace):
    if namespace is None:
        return u''
    if isinstance(namespace, int):
        return str(namespace)
    return u'|'.join(namespace)

def login(sysop = False):
    api = 'http://' + settings.host + settings.apiurl
    if sysop:
        user = (settings.sysopname, settings.sysoppwd)
    else:
        user = (settings.botname, settings.botpwd)
    return Site(api, user)

class _Namespace(object):
    MAIN, TALK, USER, USER_TALK, PROJECT, PROJECT_TALK, FILE, FILE_TALK, MEDIAWIKI, MEDIAWIKI_TALK, \
    TEMPLATE, TEMPLATE_TALK, HELP, HELP_TALK, CATEGORY, CATEGORY_TALK = xrange(16)

class Site(object):
    def __init__(self, api, user=None, **config): # put user into config?
        self._api = api
        self._cookies = CookieJar()
        if 'maxLag' in config:
            self._maxLag = config['maxLag']
        if isinstance(user, tuple) or isinstance(user, list):
            user = {'name': user[0], 'password': user[1]}
        if user:
            token = self._apiRequest(
                action = 'login',
                lgname = user['name'],
                lgpassword = user['password'],
            )['login']['token']
            result = self._apiRequest(
                action = 'login',
                lgname = user['name'],
                lgpassword = user['password'],
                lgtoken = token,
            )['login']
            if result['result'] != 'Success':
                raise LoginFailed()
        self._revisions = {} # shared revisions?
    namespace = _Namespace()
    
    def _getUser(self):
        if not hasattr(self, '_user'):
            self._user = self._apiRequest(
                action = 'query',
                meta = 'userinfo',
            )['query']['userinfo']['name']
        return self._user
    user = property(_getUser)
    
    def _request(self, data = None, url = None):
        if url is None:
            url = self._api
        if hasattr(self, '_maxLag'):
            url += '&' if '?' in url else '?'
            url += 'maxlag='
            url += str(self._maxLag)
        while True:
            opener = urllib2.build_opener(MediaWikiCookieProcessor(self._cookies))
            opener.addheaders = [('User-agent', 'MWBOT GAE Edition')]
            api = opener.open(url, data)
            if 'X-Database-Lag' not in api.headers:
                break
            retry = float(api.headers['Retry-After'])
            api.close()
            time.sleep(retry)
        result = api.read()
        api.close()
        return result
    
    def _apiRequest(self, **data):
        data['format'] = 'json'
        data = _formData(**data)
        result = json.loads(self._request(data))
        if 'error' in result:
            raise MediaWikiApiError(result['error']['code'], result['error']['info'], data)
        return result
    
    def _token(self, action, cache=True):
        # they never changes, so cache them, but is user-specific (or session-specific?)
        if not hasattr(self, '_cachedToken'):
            token = self._apiRequest(
                action = 'query',
                prop = 'info',
                titles = '-',
                intoken = 'edit', # all of them are the same
            )['query']['pages'].values()[0]['edittoken']
            self._cachedToken = token
        return self._cachedToken
    
    def getPageByTitle(self, title, isNormalized=False):
        return _Page(self, title, isNormalized)
    __call__ = getPageByTitle
    
    def getPageByLink(self, link):
        # this method can handle simp-trad problems
        for x in (u'[', u'{', u'}', u']'):
            if x in link:
                return None
        result = self._apiRequest(
            action = 'parse',
            text = u'[[%s]]' % link,
        )['parse']['links']
        if len(result) == 0:
            return None
        else:
            return self.getPageByTitle(result[0]['*'], True)
    
    def expandTemplates(self, text):
        # this method can expand templates
        result = self._apiRequest(
            action = 'expandtemplates',
            text = text
        )['expandtemplates']['*']
        return result
    
    def categoryMembers(self, cat, ns = -1, exception = ''):
        # this method can get category's members
        _r = self._apiRequest(
            action = 'query',
            list = 'categorymembers',
            cmtitle = ('Category:' + cat).replace('Category:Category:', 'Category:'),
            cmlimit = 500
        )['query']['categorymembers']
        result = []
        if isinstance(ns, int):
            ns = [ns]
        for _p in _r:
            if exception != '':
                if exception in _p['title']:
                    continue
            if ns == [-1]:
                result.append(_p['title'])
            elif _p['ns'] in ns:
                result.append(_p['title'])
        return result
    
    def categoryInfo(self, cat):
        # this method can get category infomation
        if isinstance(cat, basestring):
            cat = [cat]
        _r = self._apiRequest(
            action = 'query',
            prop = 'categoryinfo',
            titles = ('Category:' + '|Category:'.join(cat)).replace('Category:Category:', 'Category:')
        )['query']['pages']
        result = {}
        for (_i, _c) in _r.items():
            if _i < 0:
                continue
            else:
                result[_c['title']] = _c['categoryinfo']
        return result
    
    def isExists(self, titles):
        _t = ''
        if isinstance(titles, basestring):
            _t = titles
        elif isinstance(titles, list):
            _t = '|'.join(titles)
        _r = self._apiRequest(
            action = 'query',
            titles = _t
        )['query']['pages']
        result = []
        for (id, info) in _r.items():
            if int(id) > -1:
                result.append(info['title'])
        return result
    
    def _apiList(self, query, format='list', **data):
        param = dict(
            action = 'query',
        )
        param['list'] = query
        param.update(data)
        while True:
            result = self._apiRequest(**param)
            for item in result['query'][query]:
                yield item
            if 'query-continue' in result:
                cont = {}
                conts = result['query-continue'][query]
                for k in conts:
                    cont[str(k)] = conts[k] # k is unicode
                param.update(cont)
            else:
                break
    
    # we should avoid such duplicates
    def _apiProp(self, query, format='list', **data):
        param = dict(
            action = 'query',
        )
        param['prop'] = query
        param.update(data)
        while True:
            result = self._apiRequest(**param)
            page = result['query']['pages'].values()[0]
            if query not in page:
                return # no such prop items
            for item in page[query]:
                yield item
            if 'query-continue' in result:
                cont = {}
                conts = result['query-continue'][query]
                for k in conts:
                    cont[str(k)] = conts[k] # k is unicode
                param.update(cont)
            else:
                break
    
    def _getServer(self):
        if not hasattr(self, '_server'):
            # TODO: self._loadSiteInfo
            self._server = self._apiRequest(
                action = 'query',
                meta = 'siteinfo',
                siprop = 'general',
            )['query']['general']['server']
        return self._server
    server = property(_getServer)
    
    def __eq__(self, obj): # should we check usernames to be equal?
        if not isinstance(obj, self.__class__):
            return False
        return self._api == obj._api or self.server == obj.server
    
    def __hash__(self):
        return hash(self.server)

class _Page(object):
    def __init__(self, site, title, isNormalized):
        self._site = site
        self._title = title
        if isNormalized:
            self._normalizedTitle = title
    
    def move(self, target, comment=u'', moveTalk=True, moveSubpages=False, redirect=True):
        param = dict(
            action = 'move',
            to = target,
            token = self._site._token('move'),
            reason = comment,
        )
        param['from'] = self._title
        if moveTalk: param['movetalk'] = True
        if moveSubpages: param['movesubpages'] = True
        if not redirect: param['noredirect'] = True
        try:
            result = self._site._apiRequest(**param)
        except MediaWikiApiError, e:
            if e.code == 'missingtitle':
                self._title = target
                if hasattr(self, '_normalizedTitle'):
                    del self._normalizedTitle
            elif e.code != 'selfmove':
                raise PageNotMoved(e)
        else:
            self._title = target
            if hasattr(self, '_normalizedTitle'):
                del self._normalizedTitle
    
    title = property(lambda self: self._title, lambda self, target: self.move(target, moveTalk=False, moveSubpages=False, redirect=False))
    
    def _getNormalizedTitle(self):
        if not hasattr(self, '_normalizedTitle'):
            result = self._site._apiRequest(
                action = 'query',
                titles = self._title,
            )['query']
            if 'normalized' in result:
                self._normalizedTitle = result['normalized'][0]['to']
            else:
                self._normalizedTitle = self._title
        return self._normalizedTitle
    
    normalizedTitle = property(_getNormalizedTitle)
    
    def __eq__(self, obj):
        if not isinstance(obj, self.__class__):
            return False
        return self._site == obj._site and (
            self.title == obj.title or self.normalizedTitle == obj.normalizedTitle
        )
    
    def __hash__(self):
        return hash(self._site) ^ hash(self.normalizedTitle)
    
    def appendRevision(self, rev, newsection = False):
        if rev.saved:
            raise SavedRevision()
        param = dict(
            action = 'edit',
            title = self._title,
            text = rev.content,
            token = self._site._token('edit'),
            summary = rev.comment,
            recreate = True
        )
        if newsection:
            param['section'] = 'new'
        if rev.minor: param['minor'] = True
        if rev.bot: param['bot'] = True
        try:
            #print 'Appending new revision to [[%s]]...' % self._title
            result = self._site._apiRequest(**param)['edit']
        except MediaWikiApiError, e:
            raise PageNotSaved(e)
        if result['result'] != 'Success':
            raise PageNotSaved()
        if 'nochange' in result:
            return # leave revision unsaved
        rev._loadById(self._site, result['newrevid'])
    
    def _getProtection(self):
        prot = {'create': '', 'edit': '', 'move': ''}
        for p in self._site._apiRequest(
            action = 'query',
            prop = 'info',
            inprop = 'protection',
            titles = self._title,
        )['query']['pages'].values()[0]['protection']:
            prot[p['type']] = p['level']
        return prot
    protection = property(_getProtection)
    
    def _isExists(self):
        return 'missing' not in self._site._apiRequest(
            action = 'query',
            prop = 'info',
            titles = self._title,
        )['query']['pages'].values()[0]
    exists = property(_isExists)
    
    def _isPatrolled(self):
        return len(self._site._apiRequest(
            action = 'query',
            list = 'logevents',
            letype = 'patrol',
            letitle = self._title,
        )['query']['logevents']) > 0
    patrolled = property(_isPatrolled)
    
    def _getRecentRevision(self):
        result = self._site._apiRequest(
            action = 'query',
            prop = 'revisions',
            titles = self._title,
            rvprop = 'ids|flags|timestamp|user|size|comment|content',
        )['query']['pages'].values()[0]
        if 'revisions' in result:
            result = result['revisions'][0]
        else:
            raise PageNotExists()
        if result['revid'] in self._site._revisions:
            rev = self._site._revisions[result['revid']]
        else:
            rev = Revision()._loadFromResult(self._site, result)
        return rev
    recentRevision = property(_getRecentRevision)
    
    def __iadd__(self, obj):
        if isinstance(obj, Revision):
            self.appendRevision(obj)
        else:
            raise NotImplementedError()
        return self
    
    def delete(self, comment=u''):
        try:
            result = self._site._apiRequest(
                action = 'delete',
                title = self._title,
                token = self._site._token('delete'),
                reason = comment,
            )
        except MediaWikiApiError, e:
            raise PageNotDeleted(e)
        if 'delete' not in result:
            raise PageNotDeleted(result)
    
    def _linked(self, limit, namespace, viaRedirect):
        param = dict(
            bltitle = self._title,
            bllimit = limit,
            blnamespace = namespace,
        )
        if viaRedirect:
            param['blredirect'] = True
        for item in self._site._apiList(
            'backlinks',
            **param
        ):
            yield item['title']
    
    def _transcluded(self, limit, namespace):
        for item in self._site._apiList(
            'embeddedin',
            eititle = self._title,
            eilimit = limit,
            einamespace = namespace,
        ):
            yield item['title']
    
    def _members(self, limit, namespace): # TODO: .filter()
        for item in self._site._apiList(
            'categorymembers',
            cmtitle = self._title,
            cmprop = 'title',
            cmlimit = limit,
            cmnamespace = namespace,
        ):
            yield item['title']
    # members = property(lambda self: self.membersAdvanced())
    
    def _used(self, limit, namespace, viaRedirect):
        param = dict(
            iutitle = self._title,
            iulimit = limit,
            iunamespace = namespace,
        )
        if viaRedirect:
            param['iuredirect'] = True
        for item in self._site._apiList(
            'imageusage',
            **param
        ):
            yield item['title']
    
    # TODO: filterRedirects?
    def referenced(self, step=500, viaRedirect=True, namespace=None, link=False, transclude=False, member=False, use=False):
        if 1:#step <= 0:
            raise ValueError('step should be positive')
        yielded = set()
        namespace = _namespace(namespace)
        if link:
            for page in self._linked(step, namespace, viaRedirect):
                if page not in yielded:
                    yielded.add(page)
                    yield self._site.getPageByTitle(page, True)
        if transclude: # viaRedirects?
            for page in self._transcluded(step, namespace):
                if page not in yielded:
                    yielded.add(page)
                    yield self._site.getPageByTitle(page, True)
        if member: # viaRedirects?
            for page in self._members(step, namespace):
                if page not in yielded:
                    yielded.add(page)
                    yield self._site.getPageByTitle(page, True)
        if use:
            for page in self._used(step, namespace, viaRedirect):
                if page not in yielded:
                    yielded.add(page)
                    yield self._site.getPageByTitle(page, True)
    
    def _linking(self, limit, namespace):
        for item in self._site._apiProp(
            'links',
            titles = self._title,
            pllimit = limit,
            plnamespace = namespace,
        ):
            yield item['title']
    
    def _transcluding(self, limit, namespace):
        for item in self._site._apiProp(
            'templates',
            titles = self._title,
            tllimit = limit,
            tlnamespace = namespace,
        ):
            yield item['title']
    
    def _memberOf(self, limit):
        for item in self._site._apiProp(
            'categories',
            titles = self._title,
            cllimit = limit,
        ):
            yield item['title']
    
    def _using(self, limit):
        for item in self._site._apiProp(
            'images',
            titles = self._title,
            imlimit = limit,
        ):
            yield item['title']
    
    def referencing(self, step=500, namespace=None, link=False, transclude=False, member=False, use=False):
        if 1:#step <= 0:
            raise ValueError('step should be positive')
        yielded = set()
        namespace = _namespace(namespace)
        if link:
            for page in self._linking(step, namespace):
                if page not in yielded:
                    yielded.add(page)
                    yield self._site.getPageByTitle(page, True)
        if transclude:
            for page in self._transcluding(step, namespace):
                if page not in yielded:
                    yielded.add(page)
                    yield self._site.getPageByTitle(page, True)
        if member:
            for page in self._memberOf(step):
                if page not in yielded:
                    yielded.add(page)
                    yield self._site.getPageByTitle(page, True)
        if use:
            for page in self._using(step):
                if page not in yielded:
                    yielded.add(page)
                    yield self._site.getPageByTitle(page, True)

class Revision(object):
    def __init__(self, content=u'', comment=u'', minor=False, bot=False):
        # unsaved revision
        self._revid = -1
        self._minor = minor
        self._size = len(content.encode('utf8'))
        self._comment = comment
        self._content = content
        self._bot = bot
    
    def _loadById(self, site=None, revid=None):
        # lazy eval
        if site:
            if revid <= 0:
                raise ValueError('revid should > 0')
            if revid == self._revid:
                return
            if self.saved:
                raise AttributeError('Cannot set revid to another when Revision is saved')
            self._site = site
            self._revid = revid
            self._deleteAttributes('_minor', '_user', '_timestamp', '_size', '_comment', '_content')
        # if site is None, load it
        else:
            return self._loadFromResult(self._site, self._site._apiRequest(
                action = 'query',
                prop = 'revisions',
                revids = self._revid,
                rvprop = 'ids|flags|timestamp|user|size|comment|content',
            )['query']['pages'].values()[0]['revisions'][0])
    
    def _loadFromResult(self, site, result):
        if self.saved and result['revid'] != self._revid:
            raise AttributeError('Cannot set revid to another when Revision is saved')
        self._site = site
        self._revid = result['revid']
        self._minor = 'minor' in result
        self._user = result['user']
        self._timestamp = _fromTimestamp(result['timestamp'])
        self._size = result['size'] if 'size' in result else len(result['*'].encode('utf8')) # when?
        self._comment = result.get('comment', u'')
        self._content = result['*']
        self._site._revisions[self._revid] = self
        return self
    
    def _deleteAttributes(self, *attrs):
        for attr in attrs:
            if hasattr(self, attr):
                delattr(self, attr)
    
    def _getAttribute(self, attr):
        if not hasattr(self, attr):
            self._loadById()
        return getattr(self, attr)
    
    saved = property(lambda self: self._revid > 0)
    revisionId = property(lambda self: self._revid)
    minor = property(lambda self: self._getAttribute('_minor'))
    user = property(lambda self: self._getAttribute('_user'))
    timestamp = property(lambda self: self._getAttribute('_timestamp'))
    size = property(lambda self: self._getAttribute('_size'))
    comment = property(lambda self: self._getAttribute('_comment'))
    content = property(lambda self: self._getAttribute('_content'))
    bot = property(lambda self: self._bot)